# import numpy as np
#
# my_dict = {'a': 10, 'b': 20, 'c': 30}
# key='c'
# if key in my_dict:
#     value = my_dict[key]
# else:
#     value = 0
#
# print(value)

# import numpy as np
# from keras.models import Sequential
# from keras.layers import Dense
# data_size = 1
# # 搭建神经网络模型
# model = Sequential()
# model.add(Dense(units=10, input_dim=1, activation='elu'))  # 1个输入command，units个输出
# model.add(Dense(units=3)) #3个输出，控制θ1,θ2,θ3
# # 编译模型
# # optimizer = keras.optimizers.Adam(learning_rate=0.05)
# model.compile(optimizer='sgd', loss='mse')
# command = np.random.choice([0.0, 1.0], size=data_size)
# print("command=", command)
# theta_hat = model.predict(command, verbose=0)
# print("theta_hat=", theta_hat)
#
# random_rotations = np.array([np.pi/3, 0, 0])
# print("random_rotations=", random_rotations)
# loss = model.train_on_batch(command, np.array([random_rotations]))
# print(loss)

# from qiskit.primitives import Sampler
# from qiskit import QuantumCircuit
# from qiskit.circuit.library import RealAmplitudes
# # a Bell circuit
# bell = QuantumCircuit(2)
# bell.h(0)
# bell.cx(0, 1)
# bell.measure_all()
#
# # initialization of the sampler
# sampler = Sampler()
#
# # Sampler runs a job on the Bell circuit
# job = sampler.run(circuits=[bell], parameter_values=[[]], parameters=[[]])
# job_result = job.result()
# print(job_result)
# print([q.binary_probabilities() for q in job_result.quasi_dists])

# from qiskit.circuit import QuantumCircuit
# from qiskit.primitives import Estimator
# from qiskit.circuit.library import RealAmplitudes
# from qiskit.quantum_info import SparsePauliOp
# from qiskit.quantum_info import Pauli
# from qiskit.algorithms.gradients import ParamShiftEstimatorGradient
#
# # a Bell circuit
# qc = QuantumCircuit(1,1)
# qc.x(0)
# #qc.measure(0, 0)
#
# #H1 = SparsePauliOp.from_list([("Z", 1)])
# H1 = Pauli('Z')
# # 总结，三种方法都可以设置测量算符，
# # 第一种：H1 = SparsePauliOp.from_list([("Z", 1)])
# # 第二种：H1 = Pauli('Z')
# # 第三种：在run的参数直接用'Z'，job = estimator.run([qc], ['Z'])
#
# estimator = Estimator()
#
# # calculate [ <psi1(theta1)|H1|psi1(theta1)> ]
# job = estimator.run([qc], ['Z'])
# expectation_value = job.result().values
# print("expectation: ", expectation_value)
#
# gradient = ParamShiftEstimatorGradient(estimator)
# grad_result_all = gradient.run(qc, Pauli('Z'), [[]]).result().gradients # ALL (include |0> and |1>) State gradient computed with parameter shift
# print('grad_result_all  =', grad_result_all)

#---------------------------------------------------------

# 通过训练变分量子线路校正单量子比特信号（简化版），qiskit.algorithms改为安装qiskit_algorithms包
# 修改from qiskit.algorithms.gradients import ParamShiftEstimatorGradient为from qiskit_algorithms.gradients import ParamShiftEstimatorGradient

# 安装包包括：
# 1. qiskit
# 2. matplotlib
# 3. pylatexenc
# 4. seaborn
# 5. tensorflow
# 6. keras
# 7. qiskit_algorithms

import math
import numpy as np
import matplotlib.pyplot as plt
from qiskit import BasicAer, QuantumCircuit, transpile
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.circuit import Parameter
from qiskit.quantum_info import Pauli
from qiskit.primitives import Estimator
from qiskit_algorithms.gradients import ParamShiftEstimatorGradient

# 产生随机转动值（仿真对qubit的干扰噪音）
random_rotation = np.random.uniform(0, 2 * np.pi, 1)[0]
print("random_rotation=", random_rotation)
print("")

# 定义量子线路，其中rx为旋转X门，遵从右手定则（右手大拇指指向x轴方向，其他手指则为旋转方向）
qc = QuantumCircuit(1,1)
qc.rx(random_rotation, 0)  # 绕x轴转动random_rotation角度，相当于Unoisy
qc.barrier()

# 画出受干扰后qubit的布洛赫球状态
backend = BasicAer.get_backend('statevector_simulator')
psi = backend.run(transpile(qc, backend)).result().get_statevector(qc)
plot_bloch_multivector(psi)

# 打印并画出校正前的测量结果，qc_tmp为临时建立的量子线路，测量之后无实际作用
qc_tmp = qc.copy()
qc_tmp.measure(0, 0)
counts = BasicAer.get_backend('qasm_simulator').run(transpile(qc_tmp, backend)).result().get_counts(qc_tmp)
print('校正前测量结果：' , counts)
plot_histogram(counts)

# 定义变分量子线路，其中theta为控制的参数
theta = Parameter('θ')
qc.rx(theta, 0)
qc.barrier()

# 拷贝一份电路加入测量，便于显示电路全貌，无实际作用。之所以不直接在原qc上加测量，是因为后面用到Estimator计算测量结果期望和梯度，电路中不能有测量操作
qc_measure = QuantumCircuit.copy(qc)
qc_measure.measure(0, 0)
qc_measure.draw('mpl')

# 训练量子变分器
theta_c = 0  # 校准的角度参数
learning_rate = 0.05 # 设置学习率
t = 0 # 记录迭代次数
T = 10000 # 设置迭代上限
while t < T:
    # 计算测量结果期望和梯度
    estimator = Estimator()
    expectation_hat = estimator.run(qc, Pauli('Z'), [[theta_c]]).result().values[0]
    expectation_true = 1 # 真实期望值（标记）
    gradient = ParamShiftEstimatorGradient(estimator).run(qc, Pauli('Z'), [[theta_c]]).result().gradients[0][0]
    theta_c = theta_c - learning_rate * (expectation_hat - expectation_true) * gradient
    # 如果损失小于一定量则退出训练
    loss = math.pow((expectation_hat - expectation_true), 2)
    if loss < 0.001:
        break
    t = t + 1
print("量子变分器训练完成：")
print("共训练", t, "轮")
print("最终: loss =", loss)
print("theta_c =", theta_c)
print("")

# 把训练得到的校正角度带入量子电路
#bound_qc = qc.assign_parameters({theta: theta_c})
bound_qc = qc.assign_parameters({theta: theta_c})

# 画出校正后qubit的布洛赫球状态
backend = BasicAer.get_backend('statevector_simulator')
psi = backend.run(transpile(bound_qc, backend)).result().get_statevector(bound_qc)
plot_bloch_multivector(psi)

# 打印并画出校正后的测量结果
bound_qc.measure(0, 0)
counts = BasicAer.get_backend('qasm_simulator').run(transpile(bound_qc, backend)).result().get_counts(bound_qc)
print('校正后测量结果：' , counts)
plot_histogram(counts)
plt.show()
